package util.tree;

import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import dto.TreeNodeDTO;
import lombok.Data;
import org.apache.commons.lang3.StringUtils;

import java.util.*;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * @date 2021/6/29
 */
public class TreeBuilder<N, L> {

    /**
     * 初始化树
     */
    public BaseTree<N, L> initTree() {
        return new BaseTree<>();
    }

    @Data
    public static class BaseTree<N, L>
        implements TreeBuilderHelper.Leaf<TreeNodeDTO, N, L>, TreeBuilderHelper.When<TreeNodeDTO, N, L> {

        /**
         * 枝干节点
         */
        private Collection<N> nodes;
        private Function<N, TreeNodeDTO> nodeConvertFuc;
        private Collection<TreeNodeDTO> nodeList;
        /**
         * 树叶
         */
        private Collection<L> leaf;
        private Function<L, TreeNodeDTO> leafConvertFuc;
        private Collection<TreeNodeDTO> leafList;

        private Supplier<Boolean> when;
        private Collection<Predicate<TreeNodeDTO>> filters = Lists.newArrayList();

        @Override
        public TreeBuilderHelper.Node<TreeNodeDTO, N, L> gatherNode(Collection<N> nodes,
            Function<N, TreeNodeDTO> nodeConvertFuc) {
            this.nodes = Objects.isNull(nodes) || nodes.isEmpty() ? Lists.newArrayList() : nodes;
            Preconditions.checkNotNull(nodeConvertFuc);
            this.nodeConvertFuc = nodeConvertFuc;
            // nodes的结构
            return this;
        }

        @Override
        public TreeBuilderHelper.Leaf<TreeNodeDTO, N, L> gatherLeaf(Collection<L> leaf,
            Function<L, TreeNodeDTO> leafConvertFuc) {
            this.leaf = Objects.isNull(leaf) || leaf.isEmpty()  ? Lists.newArrayList() : leaf;
            Preconditions.checkNotNull(leafConvertFuc);
            this.leafConvertFuc = leafConvertFuc;
            return this;
        }

        @Override
        public TreeBuilderHelper.Node<TreeNodeDTO, N, L> filter(Predicate<TreeNodeDTO> predicate) {
            // 先保存, build的时候再执行
            filters.add(predicate);
            return this;
        }

        @Override
        public TreeBuilderHelper.When<TreeNodeDTO, N, L> when(Supplier<Boolean> when) {
            Preconditions.checkNotNull(when);
            this.when = when;
            return this;
        }

        @Override
        public TreeBuilderHelper.Node<TreeNodeDTO, N, L> thenFilter(Predicate<TreeNodeDTO> predicate) {
            Boolean condition = this.when.get();
            if (condition != null && condition) {
                filter(predicate);
            }
            this.when = null;
            return this;
        }

        @Override
        public List<TreeNodeDTO> build() {
            // 叶子结点
            leafList = Objects.isNull(leaf) || leaf.isEmpty()? Lists.newArrayList()
                : leaf.stream().map(leafConvertFuc).collect(Collectors.toList());
            // 主干节点
            nodeList = Objects.isNull(nodes) || nodes.isEmpty() ? Lists.newArrayList()
                : nodes.stream().map(nodeConvertFuc).collect(Collectors.toList());

            // 节点校验器
            Set<String> remainNodeKeySet = Sets.newHashSet();
            if (!filters.isEmpty()) {
                // 保留的node key
                // filter 叶子
                leafList = leafList.stream()
                    .filter(treeNode -> filters.stream().allMatch(predicate -> predicate.test(treeNode)))
                    .collect(Collectors.toList());
                // leaf关联的node key
                Set<String> leafParentKeySet =
                    leafList.stream().map(TreeNodeDTO::getParentKey).collect(Collectors.toSet());
                remainNodeKeySet.addAll(leafParentKeySet);

                List<TreeNodeDTO> filterNodeList = nodeList.stream()
                    .filter(treeNode -> filters.stream().allMatch(predicate -> predicate.test(treeNode)))
                    .collect(Collectors.toList());
                Set<String> directNodeKeySet =
                    filterNodeList.stream().map(TreeNodeDTO::getKey).collect(Collectors.toSet());
                remainNodeKeySet.addAll(directNodeKeySet);

                Map<String, TreeNodeDTO> nodeMap = nodeList.stream()
                    .collect(Collectors.toMap(TreeNodeDTO::getKey, Function.identity(), (o1, o2) -> o2));
                // 找到所有父节点
                Queue<String> queue = new LinkedList<>(remainNodeKeySet);
                while (!queue.isEmpty()) {
                    String key = queue.poll();
                    TreeNodeDTO treeNode = nodeMap.get(key);
                    if (treeNode != null && treeNode.getParentKey() != null) {
                        if (!remainNodeKeySet.contains(treeNode.getParentKey())) {
                            remainNodeKeySet.add(treeNode.getParentKey());
                            queue.offer(treeNode.getParentKey());
                        }
                    }
                }
            }

            nodeList = filters.isEmpty() ? nodeList : nodeList.stream()
                .filter(node -> remainNodeKeySet.contains(node.getKey())).collect(Collectors.toList());

            Map<String, List<TreeNodeDTO>> leafParentMap =
                leafList.stream().collect(Collectors.groupingBy(TreeNodeDTO::getParentKey, Collectors.toList()));
            Map<String, List<TreeNodeDTO>> nodeParentMap =
                nodeList.stream().filter(i -> Objects.nonNull(i.getParentKey()))
                    .collect(Collectors.groupingBy(TreeNodeDTO::getParentKey, Collectors.toList()));

            List<TreeNodeDTO> res = Lists.newArrayList();
            nodeList.forEach(node -> {

                List<TreeNodeDTO> leafChild = leafParentMap.getOrDefault(node.getKey(), Lists.newArrayList());
                leafChild.sort(Comparator.comparing(TreeNodeDTO::getSort));

                List<TreeNodeDTO> nodeChild = nodeParentMap.getOrDefault(node.getKey(), Lists.newArrayList());
                nodeChild.sort(Comparator.comparing(TreeNodeDTO::getSort));

                ArrayList<TreeNodeDTO> child = Lists.newArrayListWithExpectedSize(leafChild.size() + nodeChild.size());
                // 先放节点
                child.addAll(nodeChild);
                child.addAll(leafChild);
                if (Objects.isNull(child) || child.isEmpty()) {
                    // 前端大佬要求传null
                    child = null;
                }
                node.setChildren(child);
                if (root(node)) {
                    res.add(node);
                }
            });

            return res;
        }

        @Override
        public List<TreeNodeDTO> buildWithRoot(Consumer<TreeNodeDTO> rootConsumer) {
            List<TreeNodeDTO> build = build();
            build.forEach(node -> rootConsumer.accept(node));
            return build;
        }

        @Override
        public List<TreeNodeDTO> buildWithNode(Consumer<TreeNodeDTO> nodeConsumer) {
            List<TreeNodeDTO> build = build();
            nodeList.forEach(node -> nodeConsumer.accept(node));
            return build;
        }

        @Override
        public List<TreeNodeDTO> buildWithLeaf(Consumer<TreeNodeDTO> leafConsumer) {
            List<TreeNodeDTO> build = build();
            leafList.forEach(node -> leafConsumer.accept(node));
            return build;
        }

        /**
         * 是否是root
         *
         * @param treeNode
         * @return
         */
        private boolean root(TreeNodeDTO treeNode) {
            return !(Objects.nonNull(treeNode) && StringUtils.isNotBlank(treeNode.getParentKey())
                && !treeNode.getParentKey().equals("0"));
        }
    }
}
